package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.ml.linalg
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo04ImagesModel {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[*]")
      .appName("Demo04ImagesModel")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._

    val imageDF: DataFrame = spark
      .read
      .format("image")
      .load("C:\\Users\\shujia\\Desktop\\train")

    imageDF.printSchema()
    //    imageDF.show(truncate = false)

    // 1、将文件的路径及数据提取出来，并进行特征工程处理
    val newImageDF: DataFrame = imageDF.select($"image.origin" as "path", $"image.data")
      // 通过SQL的方式不好处理，那就使用类似RDD的方式对每一条数据进行处理
      .as[(String, Array[Byte])]
      .map(t2 => {
        val path: String = t2._1
        val value: Array[Byte] = t2._2
        // 将每个像素点转换成的十进制值进行处理
        // 如果大于零 将其置为0 小于0 将其置为1
        val newValue: Array[Double] = value.map(v => {
          if (v >= 0) {
            0.0
          } else {
            100.0
          }
        })
        val denseVec: linalg.Vector = Vectors.dense(newValue)

        // 从path中提取文件名
        val fileName: String = path.split("/").last
        (fileName, denseVec.toSparse)
      }).toDF("fileName", "features")

    // 加载image_res.txt数据 ，通过文件名进行关联 获取每张图片对应的数字 并作为label列
    val imageResDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", " ")
      .schema("fileName String,label Double")
      .load("spark/data/mllib/data/image_res.txt")

    val newImageResDF: DataFrame = newImageDF.join(imageResDF, List("fileName"), "inner")
      .select($"fileName", $"features", $"label")

    // 2、将数据按照8：2进行切分形成训练集、测试集
    val Array(trainDF, testDF) = newImageResDF.randomSplit(Array(0.8, 0.2))
    testDF.cache()

    // 3、选择合适的模型：逻辑回归模型
    val lr: LogisticRegression = new LogisticRegression()
      .setMaxIter(10) // 最大的迭代次数
      //      .setRegParam(0.3) // 正则化参数
      //      .setElasticNetParam(0.8)
      // 设置是否有截距
      .setFitIntercept(true)
      .setFamily("multinomial")

    // Fit the model，将训练集带入模型进行训练
    val lrModel: LogisticRegressionModel = lr.fit(trainDF)

    // 4、使用测试集对模型进行评估
    val transformDF: DataFrame = lrModel.transform(testDF)

    transformDF.show(100)
    println(s"逻辑回归模型的准确率为：${1 - transformDF.where($"label" =!= $"prediction").count() / testDF.count().toDouble}")


  }

}
